"""
使用上下文管理器的智能体示例
演示如何使用with语句自动管理智能体生命周期
"""
import json
import time
from agent_sdk import AgentSDK


def process_task(task):
    """处理任务的函数"""
    print(f"处理任务: {task['taskId']}")
    
    # 模拟处理
    time.sleep(2)
    
    return {
        "result": "处理完成",
        "timestamp": time.time()
    }


def main():
    # 使用with语句，自动注册和注销
    with AgentSDK(
        platform_url="http://localhost:8080",
        agent_config={
            "name": "ContextAgent",
            "type": "PROCESSING",
            "description": "使用上下文管理器的智能体",
            "endpoint": "http://localhost:8002"
        }
    ) as sdk:
        print(f"智能体已注册: {sdk.agent_id}")
        
        # 模拟处理一个任务
        task = {
            "taskId": "task_demo",
            "inputData": json.dumps({"data": "test"})
        }
        
        try:
            result = process_task(task)
            sdk.submit_result(
                task_id=task["taskId"],
                status="SUCCESS",
                output_data=result
            )
        except Exception as e:
            sdk.submit_result(
                task_id=task["taskId"],
                status="FAILED",
                error_message=str(e)
            )
        
        # 等待一段时间
        time.sleep(10)
    
    # 退出with块后自动注销
    print("智能体已注销")


if __name__ == "__main__":
    main()
